#!/usr/bin/env python3
import paramiko
import argparse
import json
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
import time
import yaml

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('stress_test.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)

class ChaosBladeManager:
    def __init__(self, config_file='servers.json'):
        with open(config_file, 'r', encoding='utf-8') as f:
            self.config = json.load(f)
        self.servers = self.config['servers']
        self.blade_path = self.config.get('chaosblade_path', './chaosblade-1.2.0/blade')
        self.results_cache = {}

    def execute_ssh_command(self, host, username, password, command, port=22, timeout=30):
        try:
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(host, port=port, username=username, password=password, timeout=10)

            stdin, stdout, stderr = ssh.exec_command(command, timeout=timeout)
            output = stdout.read().decode('utf-8', errors='ignore')
            error = stderr.read().decode('utf-8', errors='ignore')
            exit_code = stdout.channel.recv_exit_status()

            ssh.close()

            return {
                'host': host,
                'success': exit_code == 0,
                'output': output,
                'error': error,
                'exit_code': exit_code
            }
        except Exception as e:
            return {
                'host': host,
                'success': False,
                'output': '',
                'error': str(e),
                'exit_code': -1
            }

    def build_command(self, scenario_type, params):
        base_cmd = f"nohup {self.blade_path} create"

        if scenario_type == 'cpu':
            cpu_percent = params.get('cpu_percent', 30)
            cpu_count = params.get('cpu_count', '')
            cmd = f"{base_cmd} cpu load --cpu-percent {cpu_percent}"
            if cpu_count:
                cmd += f" --cpu-count {cpu_count}"

        elif scenario_type == 'mem':
            mem_percent = params.get('mem_percent', 50)
            mode = params.get('mode', 'ram')
            reserve = params.get('reserve', '')
            cmd = f"{base_cmd} mem load --mode {mode} --mem-percent {mem_percent}"
            if reserve:
                cmd += f" --reserve {reserve}"

        elif scenario_type == 'disk-fill':
            path = params.get('path', '/tmp')
            size = params.get('size', 1000)
            cmd = f"{base_cmd} disk fill --path {path} --size {size}"

        elif scenario_type == 'disk-burn':
            path = params.get('path', '/tmp')
            read = params.get('read', True)
            write = params.get('write', True)
            cmd = f"{base_cmd} disk burn --path {path}"
            if read:
                cmd += " --read"
            if write:
                cmd += " --write"

        elif scenario_type == 'network-delay':
            time_ms = params.get('time', 3000)
            interface = params.get('interface', 'eth0')
            offset = params.get('offset', '')
            cmd = f"{base_cmd} network delay --time {time_ms} --interface {interface}"
            if offset:
                cmd += f" --offset {offset}"

        elif scenario_type == 'network-loss':
            percent = params.get('percent', 50)
            interface = params.get('interface', 'eth0')
            cmd = f"{base_cmd} network loss --percent {percent} --interface {interface}"

        elif scenario_type == 'network-dns':
            domain = params.get('domain', 'www.example.com')
            ip = params.get('ip', '127.0.0.1')
            cmd = f"{base_cmd} network dns --domain {domain} --ip {ip}"

        elif scenario_type == 'process-kill':
            process = params.get('process', 'nginx')
            cmd = f"{base_cmd} process kill --process {process}"

        elif scenario_type == 'process-stop':
            process = params.get('process', 'nginx')
            cmd = f"{base_cmd} process stop --process {process}"

        else:
            raise ValueError(f"Unknown scenario type: {scenario_type}")

        cmd += " > /tmp/chaosblade_$(date +%s).log 2>&1 &"
        return cmd

    def execute_scenario(self, scenario_name, scenario_config, max_workers=10):
        scenario_type = scenario_config['type']
        params = scenario_config.get('params', {})
        target_servers = scenario_config.get('servers', 'all')

        if target_servers == 'all':
            servers = self.servers
        else:
            servers = [s for s in self.servers if s['host'] in target_servers]

        logging.info(f"执行场景 [{scenario_name}] 类型: {scenario_type}, 目标服务器: {len(servers)}台")

        command = self.build_command(scenario_type, params)

        results = []
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = {}
            for server in servers:
                future = executor.submit(
                    self.execute_ssh_command,
                    server['host'],
                    server['username'],
                    server['password'],
                    command,
                    server.get('port', 22)
                )
                futures[future] = server['host']

            for future in as_completed(futures):
                result = future.result()
                results.append(result)

                if result['success']:
                    logging.info(f"✓ {result['host']} - {scenario_name} 执行成功")
                else:
                    logging.error(f"✗ {result['host']} - {scenario_name} 失败: {result['error']}")

        return results

    def destroy_all(self, max_workers=10):
        logging.info(f"销毁所有ChaosBlade实验，目标服务器: {len(self.servers)}台")

        command = f"{self.blade_path} status --type create --status Active | grep -v 'COMMAND' | awk '{{print $1}}' | xargs -I {{}} {self.blade_path} destroy {{}}"

        results = []
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = {}
            for server in self.servers:
                future = executor.submit(
                    self.execute_ssh_command,
                    server['host'],
                    server['username'],
                    server['password'],
                    command,
                    server.get('port', 22)
                )
                futures[future] = server['host']

            for future in as_completed(futures):
                result = future.result()
                results.append(result)

                if result['success']:
                    logging.info(f"✓ {result['host']} - 实验销毁成功")
                else:
                    logging.warning(f"⚠ {result['host']} - {result['error']}")

        return results

    def get_status(self, max_workers=10):
        logging.info(f"查询ChaosBlade实验状态，目标服务器: {len(self.servers)}台")

        command = f"{self.blade_path} status --type create"

        results = []
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = {}
            for server in self.servers:
                future = executor.submit(
                    self.execute_ssh_command,
                    server['host'],
                    server['username'],
                    server['password'],
                    command,
                    server.get('port', 22)
                )
                futures[future] = server['host']

            for future in as_completed(futures):
                result = future.result()
                results.append(result)

                if result['success']:
                    lines = result['output'].strip().split('\n')
                    active_count = len([l for l in lines if 'Active' in l])
                    logging.info(f"{result['host']} - 活跃实验数: {active_count}")
                    if active_count > 0:
                        print(f"\n{result['host']}:\n{result['output']}")
                else:
                    logging.error(f"✗ {result['host']} - 查询失败: {result['error']}")

        return results

    def run_scenario_file(self, scenario_file, max_workers=10):
        with open(scenario_file, 'r', encoding='utf-8') as f:
            scenarios = yaml.safe_load(f)

        logging.info(f"加载场景配置文件: {scenario_file}")

        for scenario_name, scenario_config in scenarios['scenarios'].items():
            if scenario_config.get('enabled', True):
                self.execute_scenario(scenario_name, scenario_config, max_workers)

                delay = scenario_config.get('delay_after', 0)
                if delay > 0:
                    logging.info(f"等待 {delay} 秒...")
                    time.sleep(delay)
            else:
                logging.info(f"跳过已禁用的场景: {scenario_name}")

def main():
    parser = argparse.ArgumentParser(
        description='ChaosBlade批量管理工具 - 支持CPU/内存/磁盘/网络/进程故障注入',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  CPU压测:    python %(prog)s cpu --cpu-percent 40
  内存压测:   python %(prog)s mem --mem-percent 60
  磁盘填充:   python %(prog)s disk-fill --path /tmp --size 2000
  网络延迟:   python %(prog)s network-delay --time 5000 --interface eth0
  场景文件:   python %(prog)s scenario --file scenarios.yaml
  查看状态:   python %(prog)s status
  销毁所有:   python %(prog)s destroy
        """
    )

    parser.add_argument('action',
                       choices=['cpu', 'mem', 'disk-fill', 'disk-burn',
                               'network-delay', 'network-loss', 'network-dns',
                               'process-kill', 'process-stop',
                               'scenario', 'status', 'destroy'],
                       help='执行的操作类型')

    parser.add_argument('--config', default='servers.json', help='服务器配置文件')
    parser.add_argument('--workers', type=int, default=10, help='并发线程数')

    parser.add_argument('--cpu-percent', type=int, default=30, help='CPU使用率百分比')
    parser.add_argument('--cpu-count', type=int, help='占用的CPU核心数')

    parser.add_argument('--mem-percent', type=int, default=50, help='内存使用率百分比')
    parser.add_argument('--mode', default='ram', choices=['ram', 'cache'], help='内存模式')
    parser.add_argument('--reserve', type=int, help='保留内存大小(MB)')

    parser.add_argument('--path', default='/tmp', help='磁盘路径')
    parser.add_argument('--size', type=int, default=1000, help='磁盘填充大小(MB)')
    parser.add_argument('--read', action='store_true', help='磁盘读压测')
    parser.add_argument('--write', action='store_true', help='磁盘写压测')

    parser.add_argument('--time', type=int, default=3000, help='网络延迟时间(ms)')
    parser.add_argument('--offset', type=int, help='延迟偏移量(ms)')
    parser.add_argument('--interface', default='eth0', help='网络接口')
    parser.add_argument('--percent', type=int, default=50, help='丢包率百分比')

    parser.add_argument('--domain', help='DNS域名')
    parser.add_argument('--ip', help='DNS解析IP')

    parser.add_argument('--process', help='进程名称')

    parser.add_argument('--file', help='场景配置文件(YAML)')

    args = parser.parse_args()

    manager = ChaosBladeManager(args.config)

    if args.action == 'scenario':
        if not args.file:
            print("错误: 场景模式需要指定 --file 参数")
            return
        manager.run_scenario_file(args.file, args.workers)

    elif args.action == 'status':
        manager.get_status(args.workers)

    elif args.action == 'destroy':
        manager.destroy_all(args.workers)

    else:
        params = {}

        if args.action == 'cpu':
            params = {'cpu_percent': args.cpu_percent}
            if args.cpu_count:
                params['cpu_count'] = args.cpu_count

        elif args.action == 'mem':
            params = {
                'mem_percent': args.mem_percent,
                'mode': args.mode
            }
            if args.reserve:
                params['reserve'] = args.reserve

        elif args.action == 'disk-fill':
            params = {'path': args.path, 'size': args.size}

        elif args.action == 'disk-burn':
            params = {
                'path': args.path,
                'read': args.read or args.write == False,
                'write': args.write or args.read == False
            }

        elif args.action == 'network-delay':
            params = {
                'time': args.time,
                'interface': args.interface
            }
            if args.offset:
                params['offset'] = args.offset

        elif args.action == 'network-loss':
            params = {
                'percent': args.percent,
                'interface': args.interface
            }

        elif args.action == 'network-dns':
            if not args.domain or not args.ip:
                print("错误: DNS劫持需要 --domain 和 --ip 参数")
                return
            params = {'domain': args.domain, 'ip': args.ip}

        elif args.action in ['process-kill', 'process-stop']:
            if not args.process:
                print("错误: 进程操作需要 --process 参数")
                return
            params = {'process': args.process}

        scenario_config = {
            'type': args.action,
            'params': params,
            'servers': 'all'
        }

        manager.execute_scenario(args.action, scenario_config, args.workers)

if __name__ == '__main__':
    main()
